# -*- coding:utf-8 -*-
# Created by Jin(jinzhencheng@outlook.com) at 2018/08/01.

from redis import Redis
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from bd_spider.config import MySQLConfig
from bd_spider.config import RedisConfig


class RedisHelper(object):
    """
       Initialization 
    """
    def __init__(self, host=RedisConfig.HOST, port=RedisConfig.PORT, timeout=RedisConfig.TIMEOUT, password=RedisConfig.PASSWORD):
        self.redis_client = Redis(host=host,
                                  port=port,
                                  password=password,
                                  socket_timeout=timeout,
                                  db=0)

    def push_to_set(self, key, value):
        n = self.redis_client.sadd(key, value)
        return n > 0 and "success" or "error"

    def is_exists_in_set(self, key, value):
        result = self.redis_client.sismember(key, value)
        return result

    def get_set_all(self, key):
        result = self.redis_client.smembers(key)
        return result

    def pop_set(self, key):
        value = self.redis_client.spop(key)
        return value

    def publish(self, channel, message):
        result = self.redis_client.publish(channel, message)
        return result

    def subscribe(self, channel):
        pub = self.redis_client.pubsub()
        pub.subscribe(channel)
        pub.parse_response()
        return pub



class MySQLHelper(object):
    """
    Database operation class
    """
    session = None

    def __init__(self, db_name):
        uri = MySQLConfig.create_db_uri(db_name)
        self.mysql_uri = uri

    def open_driver(self):
        """
        Connection establishment
        :return:None
        """
        if not self.session:
            engine = create_engine(self.mysql_uri, pool_size=MySQLConfig.DEFAULT_POOL_MAX_SIZE)
            db_session = sessionmaker(bind=engine)
            self.session = db_session()
    pass

    def close_driver(self):
        """
        Connection be closed
        :return:None
        """
        if self.session:
            self.session.close()
            self.session = None
            """ 
            The engine will be disposed automatically, when main process finished. 
            So 'dispose' function is unnecessary.
            """
            # self.engine.dispose()
    pass


